import jieba
import random
import collections

import torch


def read_data(file_path):
    """读取数据集"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()


def preprocess_data(text, need_punct=False):
    """预处理数据集"""
    # 将大写字母变为小写
    text = text.replace('\u202f', ' ').replace('\xa0', ' ').lower()
    if need_punct:  # 使用标点符号
        def no_space(char, prev_char):
            return char in set(',.!?，。！？') and prev_char != ' '

        out = [' ' + char if i > 0 and no_space(char, text[i - 1]) else char
               for i, char in enumerate(text)]
        text = ''.join(out)
    else:
        def is_punct(char):
            return char in set(',.!?，。！？')

        out = [char if i > 0 and not is_punct(char) else '' for i, char in enumerate(text)]
        text = ''.join(out)
    return text


def split_text(text, method='word', language='eg'):
    if language == 'eg':
        return text.split(' ')
    elif language == 'ch':
        return split_ch(text, method)
    else:
        raise ValueError("不支持该语言的分词！")


def split_ch(text, method='word'):
    # 先去掉空格后，再分词
    text = text.replace(' ', '')
    if method == 'char':
        out = [char for i, char in enumerate(text)]
    elif method == 'word':
        seg_list = jieba.cut(text)
        zh_words = " ".join(seg_list)
        out = zh_words.split()
    else:
        print("没有这种分词方法！")
        raise
    return out


def tokenize_data(text, num_examples=None, method='word', language='ch'):
    """词元化数据集"""
    tokens = []
    for i, line in enumerate(text.split('\n')):
        if num_examples and i > num_examples:
            break
        tokens.append(split_text(line, method, language))
    return tokens


class Vocab:
    """文本词表"""

    # 若某个token出现次数小于min_freq次，则直接丢掉
    # reserved_tokens确定句子开始和结束的token
    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
        # 保证不报错
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []
        # 按出现频率排序
        counter = count_corpus(tokens)
        # 按照词频从大到小排列
        self._token_freqs = sorted(counter.items(), key=lambda x: x[1],
                                   reverse=True)
        # 未知词元的索引为0
        self.idx_to_token = ['<unk>'] + reserved_tokens
        self.token_to_idx = {token: idx
                             for idx, token in enumerate(self.idx_to_token)}
        for token, freq in self._token_freqs:
            if freq < min_freq:
                break
            if token not in self.token_to_idx:
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1

    def __len__(self):
        return len(self.idx_to_token)

    def __getitem__(self, tokens):
        # 若该单词未在词表中出现则返回unk
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, self.unk)
        return [self.__getitem__(token) for token in tokens]

    def to_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]

    @property
    def unk(self):  # 未知词元的索引为0
        return 0

    @property
    def token_freqs(self):
        return self._token_freqs


def count_corpus(tokens):
    """统计词元的频率,词频统计"""
    # 这里的tokens是1D列表或2D列表
    if len(tokens) == 0 or isinstance(tokens[0], list):
        # 将词元列表展平成一个列表
        tokens = [token for line in tokens for token in line]
    # 对每个token出现的次数计数
    return collections.Counter(tokens)


def load_corpus_text(file_path, max_tokens=-1):
    """返回数据集的词元索引列表和词表"""
    text = preprocess_data(read_data(file_path), need_punct=False)
    tokens = tokenize_data(text, num_examples=None, method='word', language='ch')
    vocab = Vocab(tokens)
    # 因为数据集中的每个文本行不一定是一个句子或一个段落，
    # 所以将所有文本行展平到一个列表中
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0:
        corpus = corpus[:max_tokens]
    return corpus, vocab


def seq_data_iter_random(corpus, batch_size, num_steps):
    """使用随机抽样生成一个小批量子序列"""
    # 从随机偏移量开始对序列进行分区，随机范围包括num_steps-1
    # 即获得随机开始位置
    corpus = corpus[random.randint(0, num_steps - 1):]
    # 减去1，是因为我们需要考虑标签
    num_subseqs = (len(corpus) - 1) // num_steps
    # 长度为num_steps的子序列的起始索引
    # 获得每个随机位置的下标
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # 在随机抽样的迭代过程中，
    # 来自两个相邻的、随机的、小批量中的子序列不一定在原始序列上相邻
    random.shuffle(initial_indices)

    def data(pos):
        # 返回从pos位置开始的长度为num_steps的序列
        return corpus[pos: pos + num_steps]

    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        # 在这里，initial_indices包含子序列的随机起始索引
        initial_indices_per_batch = initial_indices[i: i + batch_size]
        # 根据随机下标获得数据
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        yield torch.tensor(X), torch.tensor(Y)


def seq_data_iter_sequential(corpus, batch_size, num_steps):
    """使用顺序分区生成一个小批量子序列"""
    # 从随机偏移量开始划分序列
    offset = random.randint(0, num_steps)
    num_tokens = ((len(corpus) - offset - 1) // batch_size) * batch_size
    # 每次都是接着取数据
    Xs = torch.tensor(corpus[offset: offset + num_tokens])
    Ys = torch.tensor(corpus[offset + 1: offset + 1 + num_tokens])
    Xs, Ys = Xs.reshape(batch_size, -1), Ys.reshape(batch_size, -1)
    num_batches = Xs.shape[1] // num_steps
    for i in range(0, num_steps * num_batches, num_steps):
        X = Xs[:, i: i + num_steps]
        Y = Ys[:, i: i + num_steps]
        yield X, Y


class SeqDataLoader:
    """加载序列数据的迭代器"""

    def __init__(self, file_path, batch_size, num_steps, use_random_iter, max_tokens):
        if use_random_iter:
            self.data_iter_fn = seq_data_iter_random
        else:
            self.data_iter_fn = seq_data_iter_sequential
        self.corpus, self.vocab = load_corpus_text(file_path, max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)


def load_data_text(file_path, batch_size, num_steps, use_random_iter=False, max_tokens=10000):
    """返回指定数据集的迭代器和词表"""
    data_iter = SeqDataLoader(file_path, batch_size, num_steps, use_random_iter, max_tokens)
    return data_iter, data_iter.vocab


if __name__ == '__main__':
    file_path = "../DataSet/Xue_lyric.txt"
    text = preprocess_data(read_data(file_path), need_punct=False)
    tokens = tokenize_data(text, num_examples=None, method='word', language='ch')
    for i in range(11):
        print(tokens[i])

    vocab = Vocab(tokens)
    print(list(vocab.token_to_idx.items())[:10])
    print(vocab.idx_to_token[:10])

    for i in [3, 10]:
        print('文本:', tokens[i])
        print('索引:', vocab[tokens[i]])

    corpus, vocab = load_corpus_text(file_path)
    print(len(corpus), len(vocab))

    train_loader, data_vocab = load_data_text(file_path=file_path, batch_size=2, num_steps=5)
    print(train_loader, data_vocab)
